# -*- coding: utf-8 -*-
"""
Boundary tuner for region partition:
- KMeans init
- Boundary detection (margin + kNN mix)
- Two tuners:
  (A) incremental local search on boundary points
  (B) GA over boundary points' assignments
Evaluators:
  - proxy: fast approximate cost
  - gats : accurate cost via GA+TS (only recompute two affected regions)
"""

import math
import random
from dataclasses import dataclass
from typing import List, Tuple, Dict, Optional

import numpy as np
import pandas as pd
from sklearn.cluster import KMeans
from sklearn.neighbors import NearestNeighbors

# ========== 可选：精评估（GA+TS） ==========
# 若用 gats 评估，需要将 GA_TS.py 放在 Python 可导入路径

from GA_TS import build_instance_from_df, EvalParams, GAParams, solve, save_results_to_file
HAS_GATS = True



# ======================
# 数据载体
# ======================
@dataclass
class TaskData:
    ids: np.ndarray           # shape (N,)
    x: np.ndarray             # km
    y: np.ndarray             # km
    demand: np.ndarray        # kg
    deadline: np.ndarray      # min
    penalty: np.ndarray       # 元/分钟

    @property
    def N(self): return self.ids.shape[0]


# ======================
# KMeans 初始化
# ======================
def kmeans_init(data: TaskData, K: int, seed: int = 42) -> Tuple[np.ndarray, np.ndarray]:
    km = KMeans(n_clusters=K, n_init="auto", random_state=seed)
    coords = np.c_[data.x, data.y]
    labels = km.fit_predict(coords)
    centers = km.cluster_centers_
    return labels, centers


# ======================
# 边界点打分（margin + mix）
# ======================
def boundary_scores(coords: np.ndarray, labels: np.ndarray, centers: np.ndarray, k_nn: int = 10):
    n = coords.shape[0]
    # margin：最近与次近质心距离差距
    d_cent = np.linalg.norm(coords[:, None, :] - centers[None, :, :], axis=2)  # (n,K)
    order = np.argsort(d_cent, axis=1)
    d1 = d_cent[np.arange(n), order[:, 0]]
    d2 = d_cent[np.arange(n), order[:, 1]]
    margin = (d2 - d1) / (d2 + 1e-9)

    # mix：kNN 中不同簇比例
    nn = NearestNeighbors(n_neighbors=min(k_nn + 1, n))
    nn.fit(coords)
    _, idx = nn.kneighbors(coords)
    mix = np.zeros(n, float)
    for i in range(n):
        nbrs = idx[i][1:]
        mix[i] = np.mean(labels[nbrs] != labels[i])
    return margin, mix


def pick_boundary_indices(margin, mix, q_margin: float = 0.3, q_mix: float = 0.7):
    th_m = np.quantile(margin, q_margin)
    th_x = np.quantile(mix, q_mix)
    is_boundary = (margin <= th_m) | (mix >= th_x)
    return np.where(is_boundary)[0]


# ======================
# 评估函数
# ======================
def proxy_region_cost(idx: np.ndarray, data: TaskData, depot_xy: Tuple[float, float]) -> float:
    """
    非常快的区域代价近似：
      - 星型路径近似: 2 * sum(dist(depot, i))
      - 简化迟到罚金：按距离/速度得到到达时间顺序粗估（可按需加强）
    """
    if idx.size == 0: return 0.0
    X = data.x[idx]; Y = data.y[idx]
    dx = X - depot_xy[0]; dy = Y - depot_xy[1]
    dist = np.hypot(dx, dy)
    # 近似路长
    star_len = 2.0 * float(np.sum(dist))
    # 迟到粗估（按与仓库距离排序的顺序到达）
    order = np.argsort(dist)
    speed_kmh = 40.0
    v = speed_kmh / 60.0  # km/min
    t = 0.0
    penalty = 0.0
    for pos in order:
        i = idx[pos]
        leg = dist[pos]
        t += leg  # 去
        late = max(0.0, t - float(data.deadline[i]))
        penalty += float(data.penalty[i]) * late
        t += leg  # 回
    return star_len + penalty


def gats_region_cost(idx: np.ndarray, data: TaskData,REGION_ID, depot_xy: Tuple[float, float]) -> float:
    """
    调用 GA+TS 对单个区域精评估。
    仅在 “增量微调” 中会对两个受影响区域调用，避免过慢。
    """
    if idx.size == 0: return 0.0
    if not HAS_GATS:
        raise RuntimeError("需要 GA_TS.py 以使用 gats 评估。")

    df = pd.DataFrame({
        "x坐标__km_": data.x[idx].astype(float),
        "y坐标__km_": data.y[idx].astype(float),
        "规定配送时间__分钟_": data.deadline[idx].astype(float),
        "延迟处罚成本__元_分钟_": data.penalty[idx].astype(float),
        "载重__kg_": data.demand[idx].astype(float),
    }).reset_index(drop=True)

    ep = EvalParams(
        late_weight=1.0,
        lambda_duration=0.02,
        overtime_penalty=1e9,
        seed=123,
        max_total_load_kg=500.0,
        curb_weight_kg=100.0,
        max_payload_kg=400.0,
        energy_max_range_km=100.0,
    )
    gp = GAParams(
        pop_size=40, generations=200,
        crossover_rate=0.9, mutation_rate=0.30,
        elite_frac=0.30,
        tabu_elite_k=1, tabu_iters=12, tabu_tenure=8,
        neigh_k=10, run_tabu_every=10, neigh_window=6,
        seed=123
    )
    ins = build_instance_from_df(
        df, XCOL="x坐标__km_", YCOL="y坐标__km_", LCOL="规定配送时间__分钟_",
        PCOL="延迟处罚成本__元_分钟_", WCOL="载重__kg_",
        depot_xy=depot_xy, speed_kmh=40.0, work_budget_min= 1200
    )
    _, best_eval = solve(ins, ep, gp)
    save_results_to_file(best_eval, SHEET_NAME, REGION_ID)
    return float(best_eval.cost)


def total_cost(assignments: np.ndarray, centers: np.ndarray, data: TaskData,
               evaluator: str = "proxy") -> float:
    K = centers.shape[0]
    cost = 0.0
    for r in range(K):
        idx = np.where(assignments == r)[0]
        depot = tuple(centers[r])
        if evaluator == "proxy":
            cost += proxy_region_cost(idx, data, depot)
        elif evaluator == "gats":
            cost += gats_region_cost(idx, data, r, depot)
        else:
            raise ValueError("evaluator must be 'proxy' or 'gats'")
    return cost


# ======================
# 增量微调（仅动两个受影响区域）
# ======================
def incremental_tune(assignments: np.ndarray, centers: np.ndarray, data: TaskData,
                     boundary_idx: np.ndarray, max_iters: int = 200,
                     evaluator: str = "proxy", seed: int = 123,
                     candidate_regions_per_point: int = 2) -> Tuple[np.ndarray, float]:
    """
    对边界点做 move (r->s)，每步只重算两个区域的代价。
    evaluator:
      - 'proxy': 超快（推荐日常微调）
      - 'gats' : 精评估（仅在受影响的两个区调用 GA+TS）
    """
    random.seed(seed); np.random.seed(seed)
    K = centers.shape[0]
    coords = np.c_[data.x, data.y]
    # 为每个点预先准备“候选目标区域”（最近的若干个中心）
    d2c = np.linalg.norm(coords[:, None, :] - centers[None, :, :], axis=2)  # (N, K)
    cand_regions = np.argsort(d2c, axis=1)[:, :min(candidate_regions_per_point+1, K)]

    assign = assignments.copy()
    cur_cost = total_cost(assign, centers, data, evaluator=evaluator)

    for it in range(max_iters):
        improved = False
        # 随机化边界点顺序
        pts = boundary_idx.copy()
        np.random.shuffle(pts)

        for i in pts:
            r = int(assign[i])
            # 按与中心距离的顺序试最近的目标区
            for s in cand_regions[i]:
                s = int(s)
                if s == r: 
                    continue

                # 两区域 index
                idx_r = np.where(assign == r)[0]
                idx_s = np.where(assign == s)[0]

                # before
                depot_r = tuple(centers[r]); depot_s = tuple(centers[s])
                if evaluator == "proxy":
                    before = proxy_region_cost(idx_r, data, depot_r) + proxy_region_cost(idx_s, data, depot_s)
                else:
                    before = gats_region_cost(idx_r, data,r, depot_r) + gats_region_cost(idx_s, data, s, depot_s)

                # after：把 i 从 r 移到 s
                idx_r_new = idx_r[idx_r != i]
                idx_s_new = np.concatenate([idx_s, np.array([i], dtype=int)])
                if evaluator == "proxy":
                    after = proxy_region_cost(idx_r_new, data, depot_r) + proxy_region_cost(idx_s_new, data, depot_s)
                else:
                    after = gats_region_cost(idx_r_new, data, r, depot_r) + gats_region_cost(idx_s_new, data, s, depot_s)

                delta = after - before
                if delta < -1e-9:
                    # 接受改进
                    assign[i] = s
                    cur_cost += delta
                    improved = True
                    break  # 该点成功移动，换下一个点

        if not improved:
            break

    return assign, cur_cost


# ======================
# GA 微调（染色体仅包含“边界点→目标区域”）
# ======================
def ga_boundary_tune(assignments: np.ndarray, centers: np.ndarray, data: TaskData,
                     boundary_idx: np.ndarray, pop: int = 40, gens: int = 80,
                     pc: float = 0.9, pm: float = 0.2, elite_frac: float = 0.2,
                     evaluator: str = "proxy", seed: int = 123,
                     candidate_regions_per_point: int = 2) -> Tuple[np.ndarray, float]:
    """
    染色体：对每个边界点 i，选择其“最近的候选区域”中的一个作为归属（允许包含当前区）。
    其它非边界点保持不变；适应度=全局总成本。
    """
    random.seed(seed); np.random.seed(seed)
    K = centers.shape[0]
    coords = np.c_[data.x, data.y]
    d2c = np.linalg.norm(coords[:, None, :] - centers[None, :, :], axis=2)
    cand_regions = np.argsort(d2c, axis=1)[:, :min(candidate_regions_per_point+1, K)]

    bidx = list(boundary_idx.tolist())
    B = len(bidx)
    if B == 0:
        return assignments.copy(), total_cost(assignments, centers, data, evaluator=evaluator)

    # 基因编码：每位是 cand_regions[i] 的下标（而非直接区域号）
    def decode_gene(gene_idx: int, gene_val: int) -> int:
        i = bidx[gene_idx]
        return int(cand_regions[i, gene_val])

    def phenotype(gene: np.ndarray) -> np.ndarray:
        assign = assignments.copy()
        for gi, val in enumerate(gene):
            i = bidx[gi]
            assign[i] = decode_gene(gi, int(val))
        return assign

    def fitness(gene: np.ndarray) -> float:
        assign = phenotype(gene)
        cost = total_cost(assign, centers, data, evaluator=evaluator)
        return 1.0 / (1.0 + cost)

    # 初始化种群：包含当前解（作为精英）+ 随机
    popz = []
    # 把当前解转基因
    base_gene = np.zeros(B, dtype=int)
    for gi, i in enumerate(bidx):
        # 找到当前区域在 cand 中的位置，如果不在则设为 0
        cur_r = int(assignments[i])
        cand = cand_regions[i]
        pos = np.where(cand == cur_r)[0]
        base_gene[gi] = int(pos[0]) if pos.size > 0 else 0
    popz.append((base_gene.copy(), fitness(base_gene)))

    while len(popz) < pop:
        g = np.zeros(B, dtype=int)
        for gi, i in enumerate(bidx):
            cand_len = cand_regions[i].shape[0]
            g[gi] = np.random.randint(0, cand_len)
        popz.append((g, fitness(g)))

    elite_k = max(1, int(pop * elite_frac))

    def select():
        a, b = random.sample(popz, 2)
        return a if a[1] > b[1] else b

    def crossover(g1: np.ndarray, g2: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        if random.random() > pc:
            return g1.copy(), g2.copy()
        pt = np.random.randint(1, B)
        c1 = np.concatenate([g1[:pt], g2[pt:]])
        c2 = np.concatenate([g2[:pt], g1[pt:]])
        return c1, c2

    def mutate(g: np.ndarray):
        for i in range(B):
            if random.random() < pm:
                cand_len = cand_regions[bidx[i]].shape[0]
                g[i] = np.random.randint(0, cand_len)

    # 进化
    for _ in range(gens):
        popz.sort(key=lambda x: x[1], reverse=True)
        new_pop = popz[:elite_k]
        while len(new_pop) < pop:
            p1 = select()[0]; p2 = select()[0]
            c1, c2 = crossover(p1, p2)
            mutate(c1); mutate(c2)
            new_pop.append((c1, fitness(c1)))
            if len(new_pop) < pop:
                new_pop.append((c2, fitness(c2)))
        popz = new_pop

    # 输出最佳
    popz.sort(key=lambda x: x[1], reverse=True)
    best_gene, best_fit = popz[0]
    best_assign = phenotype(best_gene)
    best_cost = (1.0 / best_fit) - 1.0
    return best_assign, float(best_cost)


# ======================
# 主入口（示例）
# ======================
def run_boundary_tuner(data: TaskData, K: int,
                       tuner: str = "incremental",   # "incremental" 或 "ga"
                       evaluator: str = "proxy",      # "proxy" 或 "gats"
                       seed: int = 42):
    # 1) KMeans init
    labels, centers = kmeans_init(data, K, seed=seed)

    # 2) 边界点
    coords = np.c_[data.x, data.y]
    margin, mix = boundary_scores(coords, labels, centers, k_nn=10)
    boundary_idx = pick_boundary_indices(margin, mix, q_margin=0.3, q_mix=0.7)

    print(f"[init] cost={total_cost(labels, centers, data, evaluator=evaluator):.3f}, "
          f"boundary={boundary_idx.size}/{data.N}")

    # 3) 调优
    if tuner == "incremental":
        new_labels, new_cost = incremental_tune(labels, centers, data, boundary_idx,
                                                max_iters=3, evaluator=evaluator,
                                                seed=seed, candidate_regions_per_point=2)
    elif tuner == "ga":
        new_labels, new_cost = ga_boundary_tune(labels, centers, data, boundary_idx,
                                                pop=40, gens=80, pc=0.9, pm=0.2,
                                                elite_frac=0.2,
                                                evaluator=evaluator, seed=seed,
                                                candidate_regions_per_point=2)
    else:
        raise ValueError("tuner must be 'incremental' or 'ga'")

    print(f"[final] cost={new_cost:.3f}")
    return new_labels, centers, new_cost


if __name__ == "__main__":
    PATH_XLSX = "附件一_A_data.xlsx"
    SHEET_NAME = "Case1"

    # 读数据并清洗
    df = pd.read_excel(PATH_XLSX, sheet_name=SHEET_NAME).copy()
    df.columns = df.columns.str.strip().str.replace(r"[^\w]", "_", regex=True)

    XCOL = "x坐标__km_"
    YCOL = "y坐标__km_"
    LCOL = "规定配送时间__分钟_"
    PCOL = "延迟处罚成本__元_分钟_"
    WCOL = "载重__kg_"

    for col in [XCOL, YCOL, LCOL, PCOL, WCOL]:
        df[col] = pd.to_numeric(df[col], errors="coerce")
    df = df.dropna(subset=[XCOL, YCOL, LCOL, PCOL, WCOL]).reset_index(drop=True)

    # 转成 TaskData
    ids = np.arange(1, len(df) + 1)
    data = TaskData(
        ids=ids,
        x=df[XCOL].to_numpy(float),
        y=df[YCOL].to_numpy(float),
        demand=df[WCOL].to_numpy(float),
        deadline=df[LCOL].to_numpy(float),
        penalty=df[PCOL].to_numpy(float),
    )

    # 选择调优器 & 评估器
    K = 6                  # 区域数，可调
    tuner = "incremental"  # 或 "ga"
    evaluator = "gats"    # 或 "gats"

    _labels, _centers, _cost = run_boundary_tuner(
        data, K=K, tuner=tuner, evaluator=evaluator, seed=42
    )
    # 选择调优器 & 评估器
    # evaluator='proxy' 超快；若要高精请用 'gats'（需要 GA_TS.py）
    # run_boundary_tuner(data, K=6, tuner="incremental", evaluator="proxy", seed=42)
    # run_boundary_tuner(data, K=6, tuner="ga", evaluator="proxy", seed=42)
    # run_boundary_tuner(data, K=6, tuner="incremental", evaluator="gats", seed=42)  # 需 GA_TS.py
